# -*- coding: utf-8 -*-
"""
@Time    : 2024/8/28 14:02 
@Author  : ZhangShenao 
@File    : load_review.py 
@Desc    : 加载亚马逊美食评论数据集(amazon-fine-food-reviews)

数据来源: https://www.kaggle.com/snap/amazon-fine-food-reviews
"""

# 导入pandas包
# Pandas是一个用于数据处理和分析的Python库
# 它提供了DataFrame数据结构,方便进行数据的读取、处理、分析等操作
import pandas as pd
# 导入tiktoken库
# tiktoken是OpenAI开发的一个库,用于从模型生成的文本中计算token数量
import tiktoken

TOP_N: int = 5  # 设置要筛选的评论数量为1000
EMBEDDING_ENCODING: str = "cl100k_base"
MAX_TOKEN: int = 4000  # 设置最大允许的token数量为4000


def load_review_data() -> pd.DataFrame:
    """加载评论数据"""

    # 加载原始数据集
    df = pd.read_csv(filepath_or_buffer="./data/fine_food_reviews_1k.csv", index_col=0)
    df = df[["Time", "ProductId", "UserId", "Score", "Summary", "Text"]]
    df = df.dropna()

    # 将 "Summary" 和 "Text" 字段组合成新的字段 "combined"
    df["combined"] = (
            "Title: " + df.Summary.str.strip() + "; Content: " + df.Text.str.strip()
    )

    # 对DataFrame进行排序,基于"Time"列,然后选取最后的2000条评论。
    # 这个假设是,我们认为最近的评论可能更相关,因此我们将对它们进行初始筛选
    df = df.sort_values("Time").tail(TOP_N * 2)
    # 丢弃"Time"列，因为我们在这个分析中不再需要它。
    df.drop("Time", axis=1, inplace=True)
    # 从'embedding_encoding'获取编码
    encoding = tiktoken.get_encoding(EMBEDDING_ENCODING)

    # 计算每条评论的token数量
    # 我们通过使用encoding.encode方法获取每条评论的token数,然后把结果存储在新的'n_tokens'列中。
    df["n_tokens"] = df.combined.apply(lambda x: len(encoding.encode(x)))

    # 取出token数不超过MAX_TOKEN的评论,防止token过长超过LLM上限
    # 我们使用.tail方法获取token数量在允许范围内的最后TOP_N条评论。
    df = df[df.n_tokens <= MAX_TOKEN].tail(TOP_N)

    print(f"评论数据加载完成,共{len(df)}条评论")
    print(f"评论数据: \n{df["combined"]}")

    return df
